import os
# 导入LangChain相关模块用于构建问答系统
from langchain_community.document_loaders import PyPDFLoader      # PDF文档加载器
from langchain_community.vectorstores import FAISS               # 向量数据库
from langchain_community.embeddings import DashScopeEmbeddings   # 通义千问嵌入模型
from langchain_community.chat_models import ChatTongyi           # 通义千问聊天模型
from langchain_text_splitters import RecursiveCharacterTextSplitter  # 文本分割器
from langchain_classic.chains.retrieval_qa.base import RetrievalQA  # 检索问答链
from langchain_core.prompts import PromptTemplate  # 提示模板                  # 提示模板

"""
使用Qwen大模型API 分析pdf 
"""
# 1. 配置通义千问API密钥
# 设置DASHSCOPE_API_KEY环境变量，用于访问通义千问API
os.environ["DASHSCOPE_API_KEY"] = "sk-481b5344e62040fa9ea41943d582dad0"  # 替换为实际密钥（获取地址见备注）

# 2. 加载PDF文档
loader = PyPDFLoader("../example.pdf")  # 替换为你的PDF路径
documents = loader.load()

# 3. 中文优化文本分割（适配Qwen长上下文）
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,      # 每个文本块大小
    chunk_overlap=200,    # 文本块重叠大小
    separators=["\n\n", "\n", "。", "！", "？"]  # 中文分段符
)
# 分割文档为多个文本块
docs = text_splitter.split_documents(documents)

# 4. 创建向量数据库（使用通义千问嵌入模型）
embeddings = DashScopeEmbeddings(model="text-embedding-v1")  # 官方推荐嵌入模型
# 创建FAISS向量数据库
vector_db = FAISS.from_documents(docs, embeddings)

# 5. 自定义Prompt模板（控制回答风格）
template = """
请严格根据以下上下文回答问题：
{context}

问题：{question}
要求：
1. 答案不超过3句话
2. 若上下文未提及，回答"未在文档中找到相关信息"
3. 避免主观推测
"""
# 创建提示模板实例
prompt = PromptTemplate(template=template, input_variables=["context", "question"])

# 6. 构建检索问答链
# 创建检索问答链实例，使用通义千问最强模型qwen-max
qa_chain = RetrievalQA.from_chain_type(
    llm=ChatTongyi(model="qwen-max", temperature=0.1),  # 使用最强qwen-max模型，低温度值确保输出稳定性
    chain_type="stuff",                        # 链类型
    retriever=vector_db.as_retriever(search_kwargs={"k": 3}),  # 检索器，检索3个相关片段
    return_source_documents=True,              # 返回源文档
    chain_type_kwargs={"prompt": prompt}       # 注入自定义模板
)

# 7. 执行问答
question = "今年的净利润是多少？"
# 调用问答链获取答案
result = qa_chain.invoke({"query": question})

# 8. 打印结果
print(f"📌 问题: {question}")
print(f"💡 答案: {result['result']}")
print("🔍 来源文档:")
# 显示前2个相关文档片段及其页码
for i, doc in enumerate(result['source_documents'][:2]):  # 显示前2个来源
    print(f"  片段 {i+1} (P{doc.metadata['page']+1}): {doc.page_content[:100]}...")